import numpy as np
import matplotlib.pyplot as plt

import numpy as np

def gauss_seidel_poisson_solver(psi, source_term, dx, dy, boundary_values, max_iter=1000, tol=1e-6):
    """
    使用高斯-赛德尔迭代法求解二维泊松方程的离散形式，包含边界条件处理。
    
    参数:
    psi : 二维numpy数组
        初始猜测解（或前一次迭代的结果），形状为(ny, nx)，其中nx, ny分别为网格的x方向和y方向点数。
    source_term : 二维numpy数组
        源项f(x, y)的离散表示，与psi形状相同。
    dx, dy : float
        空间步长，在x轴和y轴方向上。
    boundary_values : 二维numpy数组
        边界上的值，与psi形状相同，对应边界点的值已给定。
    max_iter : int, 可选
        迭代的最大次数，默认为1000。
    tol : float, 可选
        迭代停止的容差，默认为1e-6。
    
    返回:
    psi : 二维numpy数组
        迭代求解后的势函数值。
    """
    ny, nx = psi.shape
    if psi.shape != source_term.shape or psi.shape != boundary_values.shape:
        raise ValueError("psi, source_term, 和 boundary_values 的形状必须一致")
    
    psi_new = psi.copy()
    
    for iteration in range(max_iter):
        error = 0.0
        for i in range(1, ny-1):
            for j in range(1, nx-1):
                # 更新内部点，利用已更新的相邻点值
                old_val = psi_new[i, j]
                psi_new[i, j] = (
                    (source_term[i, j] 
                     + (psi_new[i+1, j] + psi[i-1, j]) / dy**2 
                     + (psi_new[i, j+1] + psi[i, j-1]) / dx**2)
                    / (2 / dx**2 + 2 / dy**2)
                )
                # 计算误差
                error = max(error, abs(old_val - psi_new[i, j]))
                
            # 更新边界点，确保它们保持不变
            psi_new[0, :] = boundary_values[0, :]
            psi_new[-1, :] = boundary_values[-1, :]
            psi_new[:, 0] = boundary_values[:, 0]
            psi_new[:, -1] = boundary_values[:, -1]
            
        # 检查收敛性
        if error < tol:
            print(f"Converged after {iteration} iterations with error {error}.")
            break
    
    return psi_new

# 定义求解区域和参数
nx, ny = 101, 101  # 网格点数
dx, dy = 1.0 / (nx - 1), 1.0 / (ny - 1)
x = np.linspace(0, 1, nx)
y = np.linspace(0, 1, ny)
X, Y = np.meshgrid(x, y)

# 构造源项
source_term = 2 * np.pi**2 * np.sin(np.pi*X) * np.sin(np.pi*Y)

# 初始化边界值（与解析解匹配）
boundary_values = np.sin(np.pi*X) * np.sin(np.pi*Y)

# 初始化解的猜测值（可以是任意值，这里设为0）
psi_initial_guess = np.zeros((ny, nx))

# 调用高斯-赛德尔迭代法求解
psi_solution = gauss_seidel_poisson_solver(psi_initial_guess, source_term, dx, dy, boundary_values)

# 计算解析解
analytic_solution = np.sin(np.pi*X) * np.sin(np.pi*Y)

# 计算并打印最大误差
max_error = np.max(np.abs(psi_solution - analytic_solution))
print(f"最大误差: {max_error}")

# 可视化结果
plt.figure(figsize=(12, 6))
plt.subplot(1, 2, 1)
plt.imshow(psi_solution, extent=(0, 1, 0, 1), cmap='viridis', origin='lower')
plt.colorbar(label='Numerical Solution')
plt.title('Numerical Solution')

plt.subplot(1, 2, 2)
plt.imshow(analytic_solution, extent=(0, 1, 0, 1), cmap='viridis', origin='lower')
plt.colorbar(label='Analytic Solution')
plt.title('Analytic Solution')

plt.show()